// Copyright FYT 邹承甫 2022

#include "post_processor/predictor.hpp"

#include <Eigen/Core>

namespace post_processor {

Predictor::Predictor(const double &air_fraction)
    : air_fraction_(air_fraction), is_ok_(false), bullet_speed_(25) {}

void Predictor::set_bullet_speed(double bullet_speed) {
  bullet_speed_ = bullet_speed;
}

void Predictor::set_air_fraction(double air_fraction) {
  air_fraction_ = air_fraction;
}

std::tuple<float, Eigen::VectorXd> Predictor::predict(
    const Eigen::VectorXd &state) {
  is_ok_ = false;
  double height, height_predict, dh, height_iter, angle;
  height_iter = state[2];
  height = state[2];
  double distance = sqrt(state[0] * state[0] + state[1] * state[1]);
  double flying_time = 0;

  int interation = 20;
  for (int i = 0; i < interation; ++i) {
    angle = atan2(height_iter, distance);
    if (!calcBulletModel(distance, angle, height_predict, flying_time)) {
      break;
    }
    dh = height / 1000.0 - height_predict;
    height_iter += dh * 1000.0;
    if (fabs(dh) < 0.001) {
      is_ok_ = true;
      break;
    }
  }
  Eigen::VectorXd predict_state = state;
  if (is_ok_) {
    predict_state[0] += state[3] * flying_time;
    predict_state[1] += state[4] * flying_time;
    predict_state[2] += state[5] * flying_time;
  }
  return std::make_tuple(angle * 180.0f / PI, predict_state);
}

bool Predictor::calcBulletModel(const double &distance, const double &angle,
                                double &height_predict, double &flying_time) const {
  if (fabs(air_fraction_) > 1 || air_fraction_ == 0 || bullet_speed_ == 0 ||
      fabs(angle) > PI / 6.0) {
    return false;
  }
  flying_time = (exp(distance / 1000.0 * air_fraction_) - 1) / air_fraction_ /
                bullet_speed_ / cos(angle);
  height_predict = bullet_speed_ * sin(angle) * flying_time -
                   0.5 * GRAVITY_ACC * flying_time * flying_time;
  return true;
}

}  // namespace post_processor